package com.nulldev.util.web.WebSocketAPI;

import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.charset.CharacterCodingException;
import java.nio.charset.Charset;
import java.nio.charset.CharsetDecoder;
import java.nio.charset.CharsetEncoder;
import java.nio.charset.CoderResult;
import java.nio.charset.CodingErrorAction;
import java.util.ArrayList;
import java.util.List;

class MessageBuilderFactory {
	interface Builder {
		boolean appendBytes(byte[] bytes);

		WebSocketMessage toMessage();
	}

	static class BinaryBuilder implements Builder {
		private List<byte[]> pendingBytes;
		private int pendingByteCount = 0;

		BinaryBuilder() {
			pendingBytes = new ArrayList<byte[]>();
		}

		public boolean appendBytes(byte[] bytes) {
			pendingBytes.add(bytes);
			pendingByteCount += bytes.length;
			return true;
		}

		public WebSocketMessage toMessage() {
			byte[] payload = new byte[pendingByteCount];
			int offset = 0;
			for (int i = 0; i < pendingBytes.size(); ++i) {
				byte[] segment = pendingBytes.get(i);
				System.arraycopy(segment, 0, payload, offset, segment.length);
				offset += segment.length;
			}
			return new WebSocketMessage(payload);
		}
	}

	static class TextBuilder implements Builder {
		private static ThreadLocal<CharsetDecoder> localDecoder = new ThreadLocal<CharsetDecoder>() {
			@Override
			protected CharsetDecoder initialValue() {
				Charset utf8 = Charset.forName("UTF8");
				CharsetDecoder decoder = utf8.newDecoder();
				decoder.onMalformedInput(CodingErrorAction.REPORT);
				decoder.onUnmappableCharacter(CodingErrorAction.REPORT);
				return decoder;
			}
		};
		private static ThreadLocal<CharsetEncoder> localEncoder = new ThreadLocal<CharsetEncoder>() {
			@Override
			protected CharsetEncoder initialValue() {
				Charset utf8 = Charset.forName("UTF8");
				CharsetEncoder encoder = utf8.newEncoder();
				encoder.onMalformedInput(CodingErrorAction.REPORT);
				encoder.onUnmappableCharacter(CodingErrorAction.REPORT);
				return encoder;
			}
		};

		private volatile StringBuilder builder;
		private volatile ByteBuffer carryOver;

		TextBuilder() {
			builder = new StringBuilder();
		}

		public boolean appendBytes(byte[] bytes) {
			// Uncomment if you want slower but more precise decoding. Useful if you're
			// splitting multi-byte utf8 chars
			// across websocket frames
			// String nextFrame = decodeStringStreaming(bytes);
			String nextFrame = decodeString(bytes);
			if (nextFrame != null) {
				builder.append(nextFrame);
				return true;
			}
			return false;
		}

		public WebSocketMessage toMessage() {
			if (carryOver != null) {
				return null;
			}
			return new WebSocketMessage(builder.toString());
		}

		/**
		 * Quicker but less precise utf8 decoding. Does not handle characters split
		 * across websocket frames.
		 * 
		 * @param bytes Bytes representing a utf8 string
		 * @return The decoded string
		 */
		private String decodeString(byte[] bytes) {
			try {
				ByteBuffer input = ByteBuffer.wrap(bytes);
				CharBuffer buf = localDecoder.get().decode(input);
				String text = buf.toString();
				return text;
			} catch (CharacterCodingException e) {
				return null;
			}
		}

		/**
		 * Left in for reference. Less efficient, but potentially catches more errors.
		 * Behavior is largely dependent on how strict the JVM's utf8 decoder is. It is
		 * possible on some JVMs to decode a string that then throws an error when
		 * attempting to return it to bytes.
		 * 
		 * @param bytes Bytes representing a utf8 string
		 * @return The decoded string
		 */
		@SuppressWarnings("unused")
		private String decodeStringStreaming(byte[] bytes) {
			try {
				ByteBuffer input = getBuffer(bytes);
				int bufSize = (int) (input.remaining() * localDecoder.get().averageCharsPerByte());
				CharBuffer output = CharBuffer.allocate(bufSize);
				for (;;) {
					CoderResult result = localDecoder.get().decode(input, output, false);
					if (result.isError()) {
						return null;
					}
					if (result.isUnderflow()) {
						break;
					}
					if (result.isOverflow()) {
						// We need more room in our output buffer
						bufSize = 2 * bufSize + 1;
						CharBuffer o = CharBuffer.allocate(bufSize);
						output.flip();
						o.put(output);
						output = o;
					}
				}
				if (input.remaining() > 0) {
					carryOver = input;
				}
				// Re-encode to work around bugs in UTF-8 decoder
				CharBuffer test = CharBuffer.wrap(output);
				localEncoder.get().encode(test);
				output.flip();
				String text = output.toString();
				return text;
			} catch (CharacterCodingException e) {
				return null;
			}
		}

		private ByteBuffer getBuffer(byte[] bytes) {
			if (carryOver != null) {
				ByteBuffer buffer = ByteBuffer.allocate(bytes.length + carryOver.remaining());
				buffer.put(carryOver);
				carryOver = null;
				buffer.put(bytes);
				buffer.flip();
				return buffer;
			} else {
				return ByteBuffer.wrap(bytes);
			}
		}
	}

	static Builder builder(byte opcode) {
		if (opcode == WebSocket.OPCODE_BINARY) {
			return new BinaryBuilder();
		} else {
			// Text
			return new TextBuilder();
		}
	}
}
